
import { log } from "cc";
import { CCCAgentsError } from "../CCCAgentsError";
import { BinaryWriter } from "../Utils/BinaryWrier";
import { Guid } from "../Utils/Guid";
import { _types_Constructor } from "../Utils/Types";
import { SideChannel } from "./SideChannel";
import { BinaryReader } from "../Utils/BinaryReader";



export type CachedSideChannelMessage = {
    channelId: Guid;
    message: Uint8Array;
}

/**
 * @en
 * Collection of static utilities for managing the registering/unregistering of
 *  SideChannels and the sending/receiving of messages for all the channels.
 * 
 * @zh
 * 用于管理注册/注销侧通道以及发送/接收所有通道消息的静态实用程序集合。
 */
export class SideChannelManager {


    private static _registeredChannels: Map<Guid, SideChannel> = new Map<Guid, SideChannel>();

    private static _cachedMessages: CachedSideChannelMessage[] = [];

    /**
     * @en
     * Register a side channel to begin sending and receiving messages. This method is 
     *  available for environments that have custom side channels. All built-in side
     *  channels within the ML-Agents Toolkit are managed internally and do not need to
     *  be explicitly registered/unregistered. A side channel may only be registered once.
     * 
     * @zh
     * 注册一个侧通道以开始发送和接收消息。此方法适用于具有自定义侧通道的环境。
     * ML-Agents 工具包中的所有内置侧通道都是在内部管理的，不需要显式注册/注销。侧通道只能注册一次。
     * @param sideChannel 
     */
    public static registerSideChannel(sideChannel: SideChannel): void {
        const channelId = sideChannel.channelId;
        if (this._registeredChannels.has(channelId)) {
            throw new CCCAgentsError(`A side channel with the id ${channelId} is already registered`);
        }

        const numMessages = this._cachedMessages.length;

        for (let i = 0; i < numMessages; i++) {
            const cachedMessage = this._cachedMessages.shift();
            if (cachedMessage) {
                if (channelId === cachedMessage?.channelId) {
                    sideChannel.processMessage(cachedMessage.message);
                } else {
                    this._cachedMessages.push(cachedMessage);
                }
            }

        }

        this._registeredChannels.set(channelId, sideChannel);
    }

    /**
     * @en Unregister a side channel to stop sending and receiving messages. This method is
        /// available for environments that have custom side channels. All built-in side
        /// channels within the ML-Agents Toolkit are managed internally and do not need to
        /// be explicitly registered/unregistered. Unregistering a side channel that has already
        /// been unregistered (or never registered in the first place) has no negative side effects.
        /// Note that unregistering a side channel may not stop the Python side
        /// from sending messages, but it does mean that sent messages with not result in a call
        /// to <see cref="SideChannel.OnMessageReceived(IncomingMessage)"/>. Furthermore,
        /// those messages will not be buffered and will, in essence, be lost.
     * 
        @zh 注销侧通道以停止发送和接收消息。此方法适用于具有自定义侧通道的环境。 ML-Agents 工具包中的所有内置侧通道都是在内部管理的，不需要显式注册/注销。
        注销已经注销的侧通道（或者根本没有注册的侧通道）不会产生负面影响。请注意，注销侧通道可能不会阻止 Python 侧发送消息，但这意味着发送的消息不会导致调用 
        <see cref="SideChannel.OnMessageReceived(IncomingMessage)"/>。
        此外，这些消息将不会被缓冲，实质上将丢失。 
     * 
     * @param sideChannel 
     */
    public static unregisterSideChannel(sideChannel: SideChannel): void {
        if (this._registeredChannels.has(sideChannel.channelId)) {
            this._registeredChannels.delete(sideChannel.channelId);
        }
    }

    public static unregisterAllSideChannels(): void {
        this._registeredChannels.clear();
    }

    /**
     * @en
     * Returns the SideChannel of Type T if there is one registered, or null if it doesn't.
        /// If there are multiple SideChannels of the same type registered, the returned instance is arbitrary.

     * @zh 
        * 如果已注册，则返回类型 T 的 SideChannel，如果没有则返回 null。
            如果注册了多个相同类型的 SideChannel，则返回的实例是任意的。
     * 
     * @param classConstructor 
     * @returns 
     */
    public static getSideChannel<T extends SideChannel>(classConstructor: _types_Constructor): T | null {
        for (const sideChannel of this._registeredChannels.values()) {
            if (sideChannel instanceof classConstructor) {
                return sideChannel as T;
            }
        }
        return null;
    }


    /**
     * @en
     * Grabs the messages that the registered side channels will send to Python at the current step
        into a singe byte array.

        @zh 获取注册的侧通道在当前步骤将发送到 Python 的消息，将其合并为一个字节数组。
     * @param sideChannels 
     * @returns 
     */
    public static getSideChannelMessage(sideChannels?: Map<Guid, SideChannel>): Uint8Array {

        if (!sideChannels) {
            sideChannels = this._registeredChannels;
        }

        if (!this.hasOutgoingMessages(sideChannels)) {
            return new Uint8Array();
        }

        const binaryWriter = new BinaryWriter();
        for (const sideChannel of sideChannels.values()) {
            const messageList = sideChannel.messageQueue;
            for (const message of messageList) {

                const channelIdByteArray = sideChannel.channelId.toByteArray();
                binaryWriter.writeByteArray(channelIdByteArray);
                binaryWriter.writeInt32(message.length);
                binaryWriter.writeByteArray(message);
            }
            sideChannel.messageQueue = [];
        }

        return binaryWriter.toByteArray();
    }

    public static hasOutgoingMessages(sideChannels: Map<Guid, SideChannel>): boolean {
        for (const sideChannel of sideChannels.values()) {
            const messageList = sideChannel.messageQueue;
            if (messageList.length > 0) {
                return true;
            }
        }
        return false;
    }


    /**
     * @en  
     * Separates the data received from Python into individual messages for each registered side channel.
     * 
     * @zh
     * 将从 Python 接收到的数据分离为每个注册的侧通道的单独消息。
     * @param dataReceived 
     * @param sideChannels 
     * @returns 
     */
    public static processSideChannelData(dataReceived: Uint8Array, sideChannels?: Map<Guid, SideChannel>): void {
        if (!sideChannels) {
            sideChannels = this._registeredChannels;
        }

        while (this._cachedMessages.length !== 0) {
            const cachedMessage = this._cachedMessages.shift();
            if (cachedMessage) {
                if (sideChannels.has(cachedMessage.channelId)) {
                    sideChannels.get(cachedMessage.channelId)!.processMessage(cachedMessage.message);
                }
                else {
                    log(`Unknown side channel data received. Channel Id is: ${cachedMessage.channelId}`);
                }
            }
        }

        if (dataReceived.length === 0) {
            return;
        }

        const binaryReader = new BinaryReader(dataReceived.buffer);
        while (binaryReader.length > binaryReader.offset) {
            let channelId: Guid = null!;
            let message: Uint8Array = null!;
            try {
                channelId = new Guid(binaryReader.readBytes(16));
                const messageLength = binaryReader.readInt32();
                message = binaryReader.readBytes(messageLength);
            } catch (ex) {
                throw new CCCAgentsError(`There was a problem reading a message in a SideChannel. 
                Please make sure the version of MLAgents in Unity is compatible with the Python version. 
                Original error: ${ex}`);
            }

            if (sideChannels.has(channelId)) {
                sideChannels.get(channelId)!.processMessage(message);
            } else {
                // Don't recognize this ID, but cache it in case the SideChannel that can handle
                // it is registered before the next call to ProcessSideChannelData.
                this._cachedMessages.push({
                    channelId: channelId,
                    message: message
                });
            }
        }
    }

}